import torch
import time

# 定义矩阵的大小
matrix_size = 5000

# 创建两个随机矩阵
x = torch.randn(matrix_size, matrix_size)
y = torch.randn(matrix_size, matrix_size)

# CPU 计算
start_time_cpu = time.time()
# 在 CPU 上进行矩阵乘法
result_cpu = torch.matmul(x, y)
end_time_cpu = time.time()
cpu_time = end_time_cpu - start_time_cpu
print(f"CPU 计算时间: {cpu_time} 秒")

# 检查 CUDA 是否可用
if torch.cuda.is_available():
    # 将矩阵移动到 GPU 上
    x_gpu = x.cuda()
    y_gpu = y.cuda()

    # GPU 计算
    start_time_gpu = time.time()
    # 在 GPU 上进行矩阵乘法
    result_gpu = torch.matmul(x_gpu, y_gpu)
    # 同步 GPU 操作，确保计算完成
    torch.cuda.synchronize()
    end_time_gpu = time.time()
    gpu_time = end_time_gpu - start_time_gpu
    print(f"GPU 计算时间: {gpu_time} 秒")
    print(f"GPU 比 CPU 快: {cpu_time / gpu_time} 倍")
else:
    print("CUDA 不可用，无法进行 GPU 计算。")